#!/usr/bin/python
import simplejson, MuConnection, time
from StringIO import StringIO
from xml.etree import ElementTree

""" MuRestApi contains the basic Mu REST API commands necisary to execute all DoS, PVA, Studio Fuzz, and admin/maintanence functions 
	
	Classes contained: 
	MuRestApi - This class is instantiated to run Fuzz anaysis, backups, and packet capture 
	operations and recieves one argument:
		MuConnection - The connection class for maintaining a SSL session to the Mu Test Suite

	TestCase: Not implemented yet: This class will contain a single line of data to execute a test case with Studio Verify
	
	fault: this class is used to represent fault data that can be retrieved from the API
	
	TODO: data mine analysis archive for remediation data and add a archive class	
"""

class MuRestApi():
	""" MuRestApi: MuRestApi - This class is instantiated to run Fuzz anaysis, backups, and packet capture 
		operations and recieves one argument:

		Arguments: 
		MuConnection - The connection class for maintaining a SSL session to the Mu Test Suite

		Globals:
		muConnection - The SSL connection manager
		session - containes the active session id for this instance
		VERBOSE - enabel detaield printing to screen.
	
		TODO:
		logging - Used to loog activiteis instead of VERBOSE
		archive - a class that represents a mu test result archive fro data mining
	
	"""

	muConnection = None
	session = None
	VERBOSE = False
	
	def __init__(self,connection):
		""" set up the TestAPI with the passed in MuConnection """
		self.muConnection = connection
	
#---------- BACKUPS ----------------
	
	def runBackup(self):
		""" initiate a backup job on the Mu Test 
			Example URL: https://$MU_IP/api/v3/backup/run
			
			Arguments:
			None
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/backup/run'
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getBackupStatus(self, jobId):
		""" get the status a backup job on the Mu 
			Example URL: https://$MU_IP/api/v3/backup/status?jobId=1124485278
			
			Arguments:
			jobId - the jobId to get the status on
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/backup/status?jobId=' + jobId
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getBackup(self, jobId, fileName):
		""" get the completed backup job on the Mu 
			Example URL: https://$MU_IP/api/v3/backup/get?jobId=1124485278
			
			Arguments:
			jobId - the jobId to get the status on
			fileName - the string filename including path to save the backub file as
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/backup/get?jobId=' + jobId
		responseData = self.muConnection.getFile(url, fileName)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData.isCompleted()
		
	def importTemplates(self, fileName):
		""" import a template from a local file onto the Mu 
			Example URL: https://$MU_IP/api/v3/templates/import
			
			Arguments:
			fileName - the string filename including path of the templates file to post to the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/templates/import'
		responseData = self.muConnection.postFile(url, fileName)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def exportTemplates(self, fileName, uuid=None, type=None, name=None):
		""" import a template from a local file onto the Mu. One must use (uuid) or (type and name)
			Example URL: https://$MU_IP/api/v3/templates/export?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
			Example URL: https://$MU_IP/api/v3/templates/export?type=scenario&name=MyScenarioName
			
			Arguments:
			fileName - the string filename including path of the templates save retrieved template 
				from the Mu
			uuid - the uuid of the template to export
			type - the type of template to export
			name - the name of the template to export
			
			Returns:
			MuConnection.MuResponseObject
		"""
		if uuid!=None:
			url = self.muConnection.TOP_LEVEL_URL + '/api/v3/templates/export?uuid=' + uuid
			responseData = self.muConnection.getFile(url,fileName)
			self.session = responseData
		elif (type!=None)and(name!=None):
			name2 = name.replace(" ","%20")
			url = self.muConnection.TOP_LEVEL_URL + '/api/v3/templates/export?type=' + type + '&name=' + name2
			responseData = self.muConnection.getFile(url,fileName)
			self.session = responseData
		else:
			print "invalid template get"
		if self.VERBOSE:
			print responseData
		return responseData
		
	def deleteTemplates(self, uuid):
		""" delete a template from the Mu.
			Example URL: https://$MU_IP/api/v3/template/delete?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
			
			Arguments:
			uuid - the uuid of the template to export
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/template/delete?uuid=' + uuid
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def listTemplates(self,templateType):
		""" list templates on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/list
			Example URL: https://$MU_IP/api/v3/scenario/list

			Arguments:
			templateType - the type of mu template to list
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/'+templateType+'/list'
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def listTemplateTypes(self):
		""" list all template types.
			Example URL: https://$MU_IP/api/v3/templates/types

			Arguments:
			none
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/templates/types'
		responseData = self.muConnection.get(url)
		self.session = responseData
		if  not responseData.hasError():
			tmpData = responseData.getRawResponse()
			tmpData = tmpData.split("\n")
			responseData.setRetrievedObjectType(responseData.LIST)
			responseData.setRetrievedObject(tmpData)
			responseData.setCompleted(True)
		if self.VERBOSE:
			print responseData
		return responseData
		
	def runAnalysis(self, uuid):
		""" run an anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/run?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the previusly run analysis uuid or the analysis template uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		return self.controlAnalysis("run",uuid)
		
	def getAnalysisStatus(self, uuid):
		""" get the status of an anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/status?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		return self.controlAnalysis("status",uuid)
		
	def abortAnalysis(self, uuid):
		""" abort an anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/stop?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		return self.controlAnalysis("stop",uuid)
		
	def pauseAnalysis(self, uuid):
		""" pause an anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/pause?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		return self.controlAnalysis("pause",uuid)
		
	def resumeAnalysis(self, uuid):
		""" resume an anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/resume?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		return self.controlAnalysis("resume",uuid)
		
	def listAnalysis(self):
		""" list all anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/list

			Arguments:
			none
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/analysis/list'
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def deleteAnalysis(self, uuid):
		""" delete an anlysis (DoS, PVA, Fuzz tests) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/delete?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		return self.controlAnalysis("delete",uuid)
	
	def controlAnalysis(self, action, uuid):
		""" used for all analysis operations (run,stop,pause,resume,delete) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/stop?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
			Example URL: https://$MU_IP/api/v3/analysis/delete?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			action - the control action to execute: run,stop,pause,resume,delete
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/analysis/'+ action + '?uuid='+uuid
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getAnalysisFaultList(self, uuid):
		""" gat the fault list for an anlysis (running, aborted, finished, failed, paused) on the Mu.
			Example URL: https://$MU_IP/api/v3/analysis/getFaultList?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/analysis/getFaultList?uuid=' + uuid
		responseData = self.muConnection.get(url)
		self.session = responseData
		if (responseData.isCompleted()==False):
			return responseData
		myTree = ElementTree.XML(responseData.getRawResponse())
		for element in myTree:
			print element.tag
			if element.tag == "fault_list":
				childList = element.getchildren()
				faultList =[]
				for element2 in childList:
					if element2.tag == "tba_fault":
						myFault = fault(element2)
						faultList.append(myFault)
				responseData.setRetrievedObjectType(responseData.FAULT_LIST)
				responseData.setRetrievedObject(faultList)
		if self.VERBOSE:
			print responseData
		return responseData
		
	def runArchive(self, uuid, title = None, notes = None, includeEVTs = True, includePcaps = True, includeSamples = True, includeLogo = True, includeAudit = True, sign = False ):
		""" start an archie job on an analysis (running, aborted, finished, failed) on the Mu.
			Example URL: https://$MU_IP/api/v3/archive/run?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

			Arguments:
			uuid - the analysis uuid on the Mu
			title - the Title to place at the top of the anaylsis reports
			notes - string of data to add to notes field of reports
			includeEVTs - boolean to include or exclude evts, default = true
			includePcaps - boolean to include or exclude test result pcaps, default = true
			includeSamples - boolean to include or exclude sample pcaps, default = true
			includeLogo - boolean to include or exclude logo if one is loaded on the Mu, default = true
			includeAudit - boolean to include or exclude audit, default = true
			sign - boolean to indicate id digital signature is eneded, default  = false
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/archive/run?uuid=' + uuid
		if notes!=None:
			url = url + '&notes=' + notes.replace(" ","%20")
		if title!=None:
			url = url + '&title=' + title.replace(" ","%20")
		tmpS = ''
		if includeEVTs:
			tmpS = 'evts'
		if includePcaps:
			if tmpS == '':
				tmpS = 'pcaps'
			else:
				tmpS = tmpS + ',' + 'pcaps'
		if includeSamples:
			if tmpS == '':
				tmpS = 'samples'
			else:
				tmpS = tmpS + ',' + 'samples'
		if includeLogo:
			if tmpS == '':
				tmpS = 'logo'
			else:
				tmpS = tmpS + ',' + 'logo'
		if includeAudit:
			if tmpS == '':
				tmpS = 'audit'
			else:
				tmpS = tmpS + ',' + 'audit'
		if tmpS!='':
			url = url + '&include=' + tmpS
		if sign:
			url = url + '&sign'
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getArchiveStatus(self, jobId):
		""" gat the status of an archive job on the Mu.
			Example URL: https://$MU_IP/api/v3/archive/status?jobId==7405e8af

			Arguments:
			jobId - the job id of the archive
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/archive/status?jobId=' + jobId
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getArchive(self, jobId, fileName):
		""" gat the completed archive job on the Mu and save to a file.
			Example URL: https://$MU_IP/api/v3/archive/get?jobId==7405e8af

			Arguments:
			jobId - the job id of the archive
			fileName - The filename including path to save the archive data in.
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/archive/get?jobId=' + jobId
		responseData = self.muConnection.getFile(url,fileName)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def runCapture(self, port):
		""" start a packet capture job on the Mu server.
			Example URL: https://$MU_IP/api/v3/capture/run?port=b2

			Arguments:
			port - the port of the mu to begin capture on.
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/capture/run?port=' + port
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getCaptureStatus(self, jobId):
		""" gat the capture status of a job on the Mu.
			Example URL: https://$MU_IP/api/v3/capture/status?jobId=7405e8af

			Arguments:
			jobId - the jobid of the capture in question
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/capture/status?jobId=' + jobId
		responseData = self.muConnection.get(url)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
	def getCapture(self, jobId, fileName):
		""" gat the packets captured by the capture job on the Mu.
			Example URL: https://$MU_IP/api/v3/capture/get?jobId=7405e8af

			Arguments:
			jobId - the jobid of the capture in question
			fileName - The filename including path to save the packet capture data in.
			
			Returns:
			MuConnection.MuResponseObject
		"""
		url = self.muConnection.TOP_LEVEL_URL + '/api/v3/capture/get?jobId=' + jobId
		responseData = self.muConnection.getFile(url,fileName)
		self.session = responseData
		if self.VERBOSE:
			print responseData
		return responseData
		
class fault():
	""" fault:  This class is represents the data that is retreivable from the xml response of the 
		getAnalysisFaults command:

		Arguments: 
		xmlElement - the xml data from the mu that contains all fault data for an indavidual fault
	
		TODO:
	"""
	
	def __init__(self,xmlElement=None):
		if xmlElement==None:
			return
		if xmlElement.tag == "tba_fault":
			childList = xmlElement.getchildren()
			for element2 in childList:
				if element2.tag == "confidence_level":
					self.setConfidenceLevel(element2.text)
				elif element2.tag == "title":
					self.setTitle(element2.text)
				elif element2.tag == "fault_detection":
					self.setFaultDetection(element2.text)
				elif element2.tag == "fault_isolation":
					self.setFaultIsolation(element2.text)
				elif element2.tag == "attack_type":
					self.setAttackType(element2.text)
		else:
			return
	
	def __str__(self):
		s = 'Confidence Level = ' + self.confidenceLevel + '; '
		s = s + 'Title = ' + self.title + '; '
		s = s + 'Fault Isolation Method = ' + self.faultIsolation + '; '
		s = s + 'Fault Detection Method = ' + self.faultDetection + '; '
		s = s + 'Attack Type = ' + self.attackType
		return s
	
	# fault cofidence levels
	ONE = '1'
	TWO = '2'
	THREE = '3'
	FOUR = '4'
	FIVE = '5'
	# islolation methods
	VECTOR = 'vector'
	VECTOR_RANGE = 'vector range'
	VECTOR_LOOP = 'vector loop'
	VARIANT_RANGE = 'variant range'
	# detection methods
	ISTRAMENTATION = 'instramentation'
	MONITOR = 'monitor'
	CONNECTION = 'connection'
	
	
	confidenceLevel = ''
	title = ''
	faultIsolation = ''
	faultDetection = ''
	attackType = ''
	
	def getConfidenceLevel(self):
		return self.confidenceLevel
		
	def getTitle(self):
		return self.title
		
	def getFaultIsolation(self):
		return self.faultIsolation
		
	def getFaultDetection(self):
		return self.faultDetection
		
	def getAttackType(self):
		return self.attackType
		
	def setConfidenceLevel(self, obj):
		self.confidenceLevel = obj
		
	def setTitle(self, obj):
		self.title = obj
		
	def setFaultIsolation(self, obj):
		self.faultIsolation = obj
		
	def setFaultDetection(self, obj):
		self.faultDetection = obj
		
	def setAttackType(self, obj):
		self.attackType = obj
		